"""
主函数，负责调用各个模块
"""
from trade import Trade
import stock
import log
import os
import strategy as st
import visual


class Backtest(object):

    def __init__(self):
        pass
    # 计算股票组合的累计收益
    def cal_return(self, trade, strategy):
        cal_date = strategy.get_cal_date()
        return_list = []
        for date in cal_date:
            return_so_far = strategy.total_return(trade, date)
            return_list.append(return_so_far)
        return return_list
    # 运行回测
    def run(self, trade, strategy):
        # 日志初始化
        os.remove("./日志/回测日志.txt")
        # 回测开始，写入日志
        log.write_log("回测开始\n"                                    
                      "回测开始日期：{0}\n"
                      "回测结束日期：{1}\n"
                      "策略名称：{2}\n"
                      "初始资金：{3}\n"
                      .format(strategy.start, strategy.end, strategy.name, trade.money))
        end_date = strategy.end
        # 调仓的日期
        change_date = strategy.get_date_list()
        # 计算收益的日期（比调仓日期要多）
        cal_date = strategy.get_cal_date()
        for date in cal_date:
            if date in change_date:
                log.write_log("%s\n" % date)
                # print("+++++++++++++++++++++%s+++++++++++++++++++++++" % date)
                sell_dict = strategy.get_sell_dict(trade)
                for share in sell_dict.keys():
                    trade.sell(stock.Stock(share), sell_dict[share], date)
                # print(trade.hold)
                # print("自由现金：{}".format(trade.free_money))
                buy_list = strategy.get_buy_list(date)
                buy_amount = strategy.cal_target_amount(trade.free_money, buy_list, date)
                for i in range(len(buy_list)):
                    trade.buy(stock.Stock(buy_list[i]), buy_amount[i], date)
                log.write_log("本日持仓情况：\n{}\n".format(trade.get_hold(date)))
                # print(trade.hold)
                # print("自由现金：{}".format(trade.free_money))

            # print("{}:{:.2f}%".format(date, cal_return * 100))
        print("回测结束，本次回测累计收益率为：{:.2f}%".format(strategy.total_return(trade, end_date) * 100))
        log.write_log("回测结束，本次回测收益率为：{:.2f}%".format(strategy.total_return(trade, end_date)*100))


if __name__ == "__main__":
    initial_money = 10000000  # 初始资金
    start_date = "2010-01-07"  # 回测开始日期
    end_date = "2012-12-31"  # 回测结束日期
    max_stock = 7  # 最大股票持数
    sep = 10  # 调仓周期

    trade_0 = Trade(initial_money)
    strategy_0 = st.Strategy("测试", max_stock, sep, start_date, end_date)
    backtest_0 = Backtest()
    backtest_0.run(trade_0, strategy_0)
    visual.draw_total_return(strategy_0, backtest_0, trade_0)
